import torch
from torchvision import datasets, transforms
from torch.utils.data import DataLoader
import torch.nn.functional as F


transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))])
train_data = datasets.MNIST(root ='../dataset', train=True, transform=transform, download=True)
train_loader = DataLoader(dataset=train_data, batch_size=64, shuffle=True)
test_data = datasets.MNIST(root ='../dataset', train=False, transform=transform, download=True)
test_loader = DataLoader(dataset=test_data, batch_size=64, shuffle=False)

class ResBlock(torch.nn.Module):
    def __init__(self,channels):
        super(ResBlock, self).__init__()
        self.channels = channels
        self.conv1 = torch.nn.Conv2d(channels,channels,kernel_size = 3,padding = 1)
        self.conv2 = torch.nn.Conv2d(channels,channels,kernel_size = 3,padding = 1)

    def forward(self,x):
        y = F.relu(self.conv1(x))
        y = self.conv2(y)
        return F.relu(x+y)



class Model(torch.nn.Module):
    def __init__(self):
        super(Model,self).__init__()
        self.conv1 = torch.nn.Conv2d(1,16,kernel_size = 5)
        self.conv2 = torch.nn.Conv2d(16,32,kernel_size = 5)
        self.pool = torch.nn.MaxPool2d(2)

        self.l1 = torch.nn.Linear(512,256)
        self.l2=torch.nn.Linear(256,128)
        self.l3=torch.nn.Linear(128,64)
        self.l4=torch.nn.Linear(64,10)

        self.resblock1 = ResBlock(16)
        self.resblock2 = ResBlock(32)

    def forward(self,x):
        in_size = x.size(0)
        x = self.pool(F.relu(self.conv1(x)))
        x = self.resblock1(x)
        x = self.pool(F.relu(self.conv2(x)))
        x = self.resblock2(x)
        x = x.view(in_size,-1)
        x=F.relu(self.l1(x))
        x=F.relu(self.l2(x))
        x=F.relu(self.l3(x))

        return self.l4(x)
model = Model()
device = ("cuda:0" if torch.cuda.is_available() else "cpu")
model = model.to(device)
criterion = torch.nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(),lr=0.01,momentum=0.5)

def train(epoch):
    loss_runtime = 0.0
    for batch,data in enumerate(train_loader,0):
        x,y=data
        x=x.to(device)
        y=y.to(device)
        y_pred = model(x)
        loss = criterion(y_pred,y)
        loss_runtime+=loss.item()
        loss_runtime/=x.size(0)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
    if epoch%10==0:
        print("after %s epochs, loss is %.8f"%(epoch+1,loss_runtime))

def test():
    correct,total=0,0
    with torch.no_grad():
        for (x,y) in test_loader:
            x = x.to(device)
            y = y.to(device)
            y_pred = model(x)
            _,prediction = torch.max(y_pred.data,dim=1)
            correct += (prediction==y).sum().item()
            total += y.size(0)
            acc = correct/total
    print("accuracy on test set is :%5f"%acc)

if __name__ == '__main__':
    for epoch in range(100):
        train(epoch)
        if epoch%19==0:
            test()
